import ReadWrite
n = 500
N = 500

def GetUser2Users(nuser):
    user_to_users = dict()
    
    for i in all_to_all_users[nuser]:
        user_to_users[canonical_users[i[0]]] = i[1]
    return user_to_users

def CalculateRating(user):
    songs_to_rating = dict()
    nuser = canonical_users.index(user)
    user_to_users = GetUser2Users(nuser)
    for users in user_to_users:
        for song in user_to_songs[users]:
            if not song in songs_to_rating:
                songs_to_rating[song] = user_to_users[users]
            else:
                songs_to_rating[song] += user_to_users[users]
    #print songs_to_rating['SOBABRB12A6701DF4B']
    songs_to_rating = sorted(songs_to_rating.keys(), key=lambda s: songs_to_rating[s], reverse=True)
   # for thing in songs_to_rating:
    #    print songs_to_rating[50]   
    return songs_to_rating

#read the neighbour data generated in the previous step
all_to_all_users = []
for i in xrange(55):
    print i
    all_to_all_users.extend(ReadWrite.readdict('../data/Parts/partition'+str(i)+'.txt'));

#read the full data
f = open('../data/kaggle_visible_evaluation_triplets.txt', 'r')
user_to_songs = dict()
for line in f:
    user, song, _ = line.strip().split('\t')
    if user in user_to_songs:
        user_to_songs[user].add(song)
    else:
        user_to_songs[user] = set([song])
f.close()

#read the user mapping data
f = open('../data/kaggle_users.txt', 'r')
canonical_users = map(lambda line: line.strip(), f.readlines())
f.close()

f = open('../data/kaggle_visible_evaluation_triplets.txt', 'r')
song_to_count = dict()
for line in f:
    _, song, _ = line.strip().split('\t')
    if song in song_to_count:
        song_to_count[song] += 1
    else:
        song_to_count[song] = 1
f.close()
songs_ordered = sorted(song_to_count.keys(), key=lambda s: song_to_count[s], reverse=True)

f = open('../data/kaggle_songs.txt', 'r')
song_to_index = dict(map(lambda line:line.strip().split(' '),f.readlines()))
f.close()
    
songs_to_rating = CalculateRating('fd50c4007b68a3737fe052d5a4f78ce8aa117f3d')
print songs_to_rating

f = open('../data/submission.txt', 'w')
userNum = 0
for user in canonical_users:
    print usernum
    usernum += 1
    if user in user_to_songs:
        songs_to_recommend = []
        songs_to_rating = CalculateRating(user)
        for song in songs_to_rating:
            if len(songs_to_recommend) >= n:
                break
            if not song in user_to_songs[user]:
                songs_to_recommend.append(song)
        for song in songs_ordered:
            if len(songs_to_recommend) >= N:
                break
            if not song in user_to_songs[user]:
                if not song in songs_to_recommend:
                    songs_to_recommend.append(song)
    else:
        songs_to_recommend = []
        for song in songs_ordered:
            if len(songs_to_recommend) >= N:
                break
            songs_to_recommend.append(song)
    # Transform song IDs to song indexes
    indices = map(lambda s: song_to_index[s],songs_to_recommend)
    # Write line for that user
    f.write(' '.join(indices) + '\n')
f.close()
print "Finished"